import os
import time

import cv2
import numpy as np
from selenium.webdriver.remote.webdriver import WebDriver
from selenium.webdriver.remote.webelement import WebElement

from libs.movement import Movement


class ImageTool:
    def __init__(self, driver: WebDriver):
        self.driver = driver
        if self.driver:
            self.driver.set_script_timeout(20)
        self.movement = Movement(driver)

    def find_target(self, target_file_name, threshold=0.8, base_elem: WebElement = None):
        if base_elem is None:
            screenshot_data = None
            for _ in range(3):
                try:
                    screenshot_data = self.driver.get_screenshot_as_png()
                    break
                except Exception as e:
                    print(e)
                    continue
            screenshot = cv2.imdecode(np.frombuffer(screenshot_data, np.uint8), -1)
        else:
            screenshot_data = base_elem.screenshot_as_png
            screenshot = cv2.imdecode(np.frombuffer(screenshot_data, np.uint8), -1)

        # cv2.imshow('Preview Image', screenshot)
        # cv2.waitKey(0)
        # cv2.destroyAllWindows()

        current_dir = os.path.dirname(os.path.abspath(__file__))
        parent_dir = os.path.dirname(current_dir)
        target_path = os.path.join(parent_dir, 'images', target_file_name)
        target = cv2.imread(target_path)
        result = cv2.matchTemplate(
            cv2.cvtColor(screenshot, cv2.COLOR_RGB2GRAY),
            cv2.cvtColor(target, cv2.COLOR_RGB2GRAY),
            cv2.TM_CCOEFF_NORMED
        )
        _, max_val, _, max_loc = cv2.minMaxLoc(result)
        if max_val > threshold:
            return max_loc
        return None
    def find_nearest_target(self, target_file_name, threshold=0.8, near_x=None, near_y=None):
        screenshot_data = self.driver.get_screenshot_as_png()
        screenshot = cv2.imdecode(np.frombuffer(screenshot_data, np.uint8), -1)
        current_dir = os.path.dirname(os.path.abspath(__file__))
        parent_dir = os.path.dirname(current_dir)
        target_path = os.path.join(parent_dir, 'images', target_file_name)
        target = cv2.imread(target_path)
        result = cv2.matchTemplate(
            cv2.cvtColor(screenshot, cv2.COLOR_RGB2GRAY),
            cv2.cvtColor(target, cv2.COLOR_RGB2GRAY),
            cv2.TM_CCOEFF_NORMED
        )

        # Find all locations where the matching value is greater than the threshold
        locations = np.where(result >= threshold)
        locations = list(zip(*locations[::-1]))  # Switch from (row, col) to (x, y) and convert to list of tuples

        if not locations:
            return None

        if near_x is not None or near_y is not None:
            # Calculate distances from (near_x, near_y) if provided and find the minimum distance
            distances = []
            for (x, y) in locations:
                if near_x is not None and near_y is not None:
                    distance = (x - near_x) ** 2 + (y - near_y) ** 2
                elif near_x is not None:
                    distance = abs(x - near_x)
                else:
                    distance = abs(y - near_y)
                distances.append((distance, (x, y)))
            _, closest_loc = min(distances, key=lambda t: t[0])
        else:
            # If neither near_x nor near_y are provided, return the first matched location
            closest_loc = locations[0]

        return closest_loc

    # 按顺序找到第一个目标
    def find_targets(self, target_files):
        loc = None
        for target_file_name in target_files:
            screenshot_data = self.driver.get_screenshot_as_png()
            screenshot = cv2.imdecode(np.frombuffer(screenshot_data, np.uint8), -1)
            current_dir = os.path.dirname(os.path.abspath(__file__))
            parent_dir = os.path.dirname(current_dir)
            target_path = os.path.join(parent_dir, 'images', target_file_name)
            target = cv2.imread(target_path)
            result = cv2.matchTemplate(
                cv2.cvtColor(screenshot, cv2.COLOR_RGB2BGR),
                cv2.cvtColor(target, cv2.COLOR_RGB2BGR),
                cv2.TM_CCOEFF_NORMED
            )
            _, max_val, _, max_loc = cv2.minMaxLoc(result)
            threshold = 0.8
            if max_val > threshold:
                loc = max_loc
                break
        return loc

    def find_all_targets(self, target_file_name, region=None, threshold=0.8, distance_threshold=50):
        screenshot_data = None
        for _ in range(3):
            try:
                screenshot_data = self.driver.get_screenshot_as_png()
                break
            except Exception as e:
                print(e)
                continue
        screenshot = cv2.imdecode(np.frombuffer(screenshot_data, np.uint8), -1)
        start_x = 0
        start_y = 0

        if region is not None:
            start_x, start_y, end_x, end_y = region
            screenshot = screenshot[start_y:end_y, start_x:end_x]

        current_dir = os.path.dirname(os.path.abspath(__file__))
        parent_dir = os.path.dirname(current_dir)
        target_path = os.path.join(parent_dir, 'images', target_file_name)
        target = cv2.imread(target_path)
        result = cv2.matchTemplate(
            cv2.cvtColor(screenshot, cv2.COLOR_RGB2GRAY),
            cv2.cvtColor(target, cv2.COLOR_RGB2GRAY),
            cv2.TM_CCOEFF_NORMED
        )
        locations = np.where(result >= threshold)
        matched_locations = []
        for pt in zip(*locations[::-1]):
            match_x = pt[0]
            match_y = pt[1]
            too_close = False
            for (mx, my) in matched_locations:
                dist = np.sqrt(abs((match_x - mx)**2 + (match_y - my)**2))
                if dist < distance_threshold:
                    too_close = True
                    break
            if too_close:
                continue
            matched_locations.append((match_x, match_y))

        # cv2.imshow('Target', target)
        # cv2.waitKey(0)
        # for loc in matched_locations:
        #     cv2.rectangle(screenshot, loc, (loc[0] + 20, loc[1] + 20), (0, 255, 0), 2)
        # cv2.imshow('Screenshot', screenshot)
        # cv2.waitKey(0)
        # cv2.destroyAllWindows()

        fixed_locations = []
        for loc in matched_locations:
            x, y = loc
            fixed_locations.append((x + start_x, y + start_y))
        return fixed_locations


    def move_to_match_target(self, target_file_name, match_x, match_y, max_delta=10):
        for _ in range(100):
            time.sleep(0.5)
            target_loc = self.find_target(target_file_name)
            # print(target_loc)
            if target_loc is None:
                break
            target_x, target_y = target_loc
            delta_x = 0
            delta_y = 0
            if match_x is not None:
                delta_x = match_x - target_x
            if match_y is not None:
                delta_y = match_y - target_y
            if abs(delta_x) <= max_delta and abs(delta_y) <= max_delta:
                break
            if delta_y >= max_delta:
                self.movement.move_up(1, self.calc_move_seconds(delta_y))
                continue
            if delta_y <= -max_delta:
                self.movement.move_down(1, self.calc_move_seconds(delta_y))
                continue
            if delta_x >= max_delta:
                self.movement.move_left(1, self.calc_move_seconds(delta_x))
                continue
            if delta_x <= -max_delta:
                self.movement.move_right(1, self.calc_move_seconds(delta_x))
                continue

    def calc_move_seconds(self, delta):
        move_seconds = min(0.1, max(abs(delta) / 1000, 0.02))
        return move_seconds
